import os
import sys
import random
import numpy as np
import xml.etree.ElementTree as ET
from xml.etree.ElementTree import Element

# 添加父目录到路径，以便导入simulation_constants
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'main'))
from simulation_constants import scenario_map2paraName

# SUMO相关导入
try:
    import traci
    import sumolib
    from sumolib import checkBinary
except ImportError:
    print("警告: SUMO库未安装或未正确配置。请确保SUMO已安装并且SUMO_HOME环境变量已设置。")
    print("可以通过以下方式安装SUMO:")
    print("1. 访问 https://eclipse.org/sumo/ 下载安装")
    print("2. 设置环境变量 SUMO_HOME 指向SUMO安装目录")
    print("3. 将 SUMO_HOME/tools 添加到Python路径")
    traci = None
    sumolib = None
    checkBinary = None

 
def traCI_SUMO2XML(showGUI=False,                   
                   sumocfg_url="",                  
                   message_gen_poss_lambda=0.15,
                   scenarios_path="",
                   index=0) -> None:
    """ 用于调用sumocfg生成同交通需求的xml文件

    args:
        showGUI (bool): 是否通过UI展示生成的场景
        sumocfg_url (str): SUMO生成场景的cfg文件路径
        message_gen_poss_lambda (float): 任务车辆通过泊松过程生成计算任务的速率参数
        scenarios_path (str): 场景文件路径
        index (int): 场景名称

    return:
        None

    """
    # 检查SUMO是否可用
    if traci is None or sumolib is None or checkBinary is None:
        print("错误: SUMO库不可用，无法生成场景文件。")
        print("请确保SUMO已正确安装并配置。")
        return
    def __indent(elem, level=0):
        """ 用于处理xml文件输出格式的子函数 """
        i = "\n" + level * "\t"
        if len(elem):
            if not elem.text or not elem.text.strip():
                elem.text = i + "\t"
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
            for elem in elem:
                __indent(elem, level + 1)
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
        else:
            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = i

    def normal_int_in_range(min_val, max_val, mu=None, sigma=2):
        """ 生成正态分布的整数，确保在指定范围内。
        :param min_val (int): 允许的最小值(包含)
        :param max_val (int): 允许的最大值(包含)
        :param mu (float, 可选): 正态分布的均值. 如果为None, 默认为区间中点
        :param sigma (float, 可选): 正态分布的标准差. 默认值为2
        :return int: 生成的符合要求的整数
        """
        # 默认均值为区间中点
        if mu is None:
            mu = (min_val + max_val) / 2
        # 生成正态分布小数
        num = random.gauss(mu, sigma)
        # 截断到 [min_val, max_val] 区间
        clamped = max(min_val, min(num, max_val))
        # 四舍五入为整数
        return round(clamped)

    if sumocfg_url == "":
        print("未输入有效sumocfg启动文件")
        return
    if scenarios_path == "":
        print("未指定场景生成目录")
        return

    root = ET.Element('root', {'id': '动态交通场景_1200m_200s'})
    tree = ET.ElementTree(root)

    # 是否显示UI界面
    if showGUI:
        sumoBinary = checkBinary('sumo-gui')
    else:
        sumoBinary = checkBinary('sumo')

    free_port = sumolib.miscutils.getFreeSocketPort()

    traci.start(cmd=[sumoBinary, "--start", "--delay", "0", "-c", sumocfg_url, "--quit-on-end", "--no-step-log", "-W", "--duration-log.disable"],
                port=free_port,
                numRetries=200,
                stdout=open(os.devnull, "w"))
    # , "--no-step-log", "-W", "--duration-log.disable"

    for step in range(0, 2000):
        if step >= 1000:
            frame = Element("frame", {'id': str(step),
                                      'num_of_vehicles': str(len(traci.vehicle.getIDList()))})
            for veh_id in traci.vehicle.getIDList():
                vehicle = Element("vehicle", {'id': veh_id,
                                              'x': str(traci.vehicle.getPosition(veh_id)[0]),
                                              'y': str(traci.vehicle.getPosition(veh_id)[1]),
                                              'speed': str(traci.vehicle.getSpeed(veh_id)),
                                              'lane': str(traci.vehicle.getLaneID(veh_id)),
                                              'angle': str(traci.vehicle.getAngle(veh_id))
                                              })
                t = np.random.poisson(lam=message_gen_poss_lambda,
                                      size=None)
                for _ in range(t):
                    data_size = max(0.75, min(2.25, np.random.normal(1.5, 0.25)))  # 正态分布, 单位为 M Byte
                    CVR = np.random.uniform(100, 500)                           # 计算转换率, 
                    componet = np.random.uniform(800, 1000)
                    task = Element("task", {'data_size': str(data_size),
                                            'computing_power_demand': str(data_size * CVR),
                                            'type': str(normal_int_in_range(1, 10)),
                                            'maximum_allowable_delay': str(data_size * componet)
                                            })
                    vehicle.append(task)
                frame.append(vehicle)
            root.append(frame)
        traci.simulationStep(step / 10)
    traci.close()
    __indent(root)
    print(index)
    tree.write(f"{scenarios_path}/{index}.xml",
               encoding='utf-8',
               xml_declaration=True)


def genScenarios(scenario_map,
                 scenario_gen_para,
                 message_gen_poss_lambda,
                 scenarios_count) -> None:
    """ 为同一场景参数生成固定数量的动态场景

    args:
        scenario_map (str): 场景类型
        scenario_gen_para (float): 场景参数车流量
        message_gen_poss_lambda (float): 任务生成速率参数
        num_scenarios (int): 场景数量

    return:
        None

    """
    scenarios_path = f"./scenario_xml/scenario_{scenario_map}_{scenario_map2paraName[scenario_map][0]}={scenario_gen_para}_lam={message_gen_poss_lambda}"
    # 如果已经为相同参数生成过场景, 则直接返回
    if os.path.exists(f"{scenarios_path}"):
        return
    os.makedirs(f"{scenarios_path}", exist_ok=False)
    sumo_cfg = f"../scenario_xml_generate/scenario_sumo_files/{scenario_map2paraName[scenario_map][1]}"
    for _ in range(scenarios_count):
        traCI_SUMO2XML(sumocfg_url=sumo_cfg,
                       message_gen_poss_lambda=message_gen_poss_lambda,
                       scenarios_path=scenarios_path,
                       index=_)


if __name__ == "__main__":
    genScenarios(scenario_map="highway",
                 scenario_gen_para=500,
                 message_gen_poss_lambda=0.5,
                 scenarios_count=10)
